from torch import nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.bias = True
        self.conv36_64 = nn.Conv2d(in_channels=12, out_channels=36, kernel_size=3, padding=1, bias=self.bias,
                                   device="cuda:0")
        self.conv64_128 = nn.Conv2d(in_channels=36, out_channels=64, kernel_size=3, padding=1, bias=self.bias,
                                    device="cuda:0")
        self.conv128_64 = nn.Conv2d(in_channels=64, out_channels=36, kernel_size=3, padding=1, bias=self.bias,
                                    device="cuda:0")
        self.conv64_36 = nn.Conv2d(in_channels=36, out_channels=12, kernel_size=3, padding=1, bias=self.bias,
                                   device="cuda:0")
        self.conv36_3 = nn.Conv2d(in_channels=12, out_channels=1, kernel_size=3, padding=1, bias=self.bias,
                                  device="cuda:0")
        self.relu = nn.ReLU()

    def forward(self, input):
        skip = input
        output1 = self.conv36_64(input)  # 得到64通道
        output1 = self.relu(output1)

        skip2 = output1
        output2 = self.conv64_128(output1)  # 得到128通道

        output2 = self.relu(output2)

        output3 = self.conv128_64(output2)  # 得到64通道
        output3 = self.relu(output3+skip2)

        output4 = self.conv64_36(output3)  # 得到32通道
        output4 = self.relu(output4+skip)
        output5 = self.conv36_3(output4)
        return output5
